import argparse
import os
import sys
import threading
import queue
import random
import sqlite3
from locust import HttpUser, task, between, events, LoadTestShape
from locust.exception import StopUser
from utils.performance_tool.config_loader import load_config
from utils.performance_tool.load_shapes import SHAPE_MAP
from utils.request_tool.send_request import RequestSender
from utils.time_tool.time_control import parse_duration
from utils.file_handling.excel_handling import ExcelUtil
from utils.interface_auth.auth_tool import Authentication
from config.all_paths import get_root_path

# ========== 参数解析 ==========
def parse_custom_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', required=True)
    parser.add_argument('--role', required=True)
    parser.add_argument('--profile_name', required=True)
    args, remaining = parser.parse_known_args()
    sys.argv = sys.argv[:1] + remaining
    return args

args = parse_custom_args()
ENV = args.env
ROLE = args.role
PROFILE_NAME = args.profile_name

# ========== 配置加载 ==========
config = load_config(env=ENV, role=ROLE)
env_config = config['env']
profile_config = config['profiles'].get(PROFILE_NAME, {})

# ========== 数据加载 ==========
data_path = os.path.join(get_root_path(), '../data/case_excel/wutong/uat/医生账号相关接口.xlsx')
testData = ExcelUtil(data_path).read_excel()

# ========== 多账号池管理 ==========
account_pool = queue.Queue()
user_list = env_config.get("account_pool", [])
for account in user_list:
    account_pool.put(account)

def get_account():
    try:
        return account_pool.get(timeout=2)
    except queue.Empty:
        return None

def return_account(account):
    account_pool.put(account)

# ========== 数据库存储 ==========
conn = sqlite3.connect("locust_results.db", check_same_thread=False)
cursor = conn.cursor()
cursor.execute('''
    CREATE TABLE IF NOT EXISTS results (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        user TEXT,
        success INTEGER,
        response_time REAL,
        timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
    )
''')
conn.commit()

def save_result(user, success, response_time):
    cursor.execute("INSERT INTO results (user, success, response_time) VALUES (?, ?, ?)",
                   (user, success, response_time))
    conn.commit()

# ========== 用户类 ==========
class BaseUser(HttpUser):
    wait_time = between(0.5, 2)
    abstract = True

    def on_start(self):
        account = get_account()
        if not account:
            print("没有可用账号")
            raise StopUser()
        self.account = account
        auth = Authentication(account["username"], account["password"], env_config.get("url"))
        method_name = {"test_wt": "auth_wutong", "test_th": "auth_th", "test_inter": "auth_inter"}.get(ENV, "auth_wutong")
        self.s = getattr(auth, method_name)()
        self.username = account["username"]

    def on_stop(self):
        return_account(self.account)

class BaselineUser(BaseUser):
    @task
    def run(self):
        r = RequestSender().send_request(self.s, testData[0])
        save_result(self.username, r.success, r.response_time)

class StressUser(BaseUser):
    @task(3)
    def run(self):
        r = RequestSender().send_request(self.s, testData[0])
        save_result(self.username, r.success, r.response_time)

class ConcurrencyUser(BaseUser):
    barrier = threading.Barrier(parties=profile_config.get("users", 1))

    @task
    def run(self):
        try:
            ConcurrencyUser.barrier.wait()
            r = RequestSender().send_request(self.s, testData[0])
            save_result(self.username, r.success, r.response_time)
        except Exception as e:
            print(f"[并发异常] {e}")
            raise StopUser()

# ========== 智能控制器 ==========
class SmartLoadShape(LoadTestShape):
    def tick(self):
        run_time = self.get_run_time()
        if run_time < 60:
            return (10, 5)
        elif run_time < 180:
            return (50, 10)
        elif run_time < 300:
            return (200, 20)

        return None

# ========== 动态装配 ==========
user_class = profile_config.get("user_class", "BaselineUser")
class TestScenario(eval(user_class)):
    pass

shape_class = profile_config.get("load_shape", "constant")
class CustomLoadShape(SHAPE_MAP.get(shape_class, SmartLoadShape)):
    shape_users = profile_config.get("users", 1)
    spawn_rate = profile_config.get("spawn_rate", 1)
    shape_duration = parse_duration(profile_config.get("duration", "1m"))
    shape_stages = profile_config.get("stages", [])

# ========== 报告统计 ==========
@events.test_stop.add_listener
def report_summary(environment, **kwargs):
    print("\n===== 测试统计报告 =====")
    print(f"总请求: {environment.stats.total.num_requests}")
    print(f"失败率: {environment.stats.total.fail_ratio:.2%}")
    print(f"平均响应时间: {environment.stats.total.avg_response_time:.2f} ms")
    print("=======================")
